import tensorflow as tf
import pandas as pd
import math

class MyModel(tf.keras.Model):
    def __init__(self, num_classes):
        super(MyModel, self).__init__()
        self.conv1 = tf.keras.layers.Conv2D(filters=8, kernel_size=(1,3), strides=(1,1), padding='same', activation='relu')
        self.conv2 = tf.keras.layers.Conv2D(filters=16, kernel_size=(1,3), strides=(1,1), padding='same', activation='relu')
        # self.lstm = tf.keras.layers.LSTM(units=16, activation='relu', kernel_regularizer=kernel_regularizer)
        self.flatten = tf.keras.layers.Flatten()
        self.fc1 = tf.keras.layers.Dense(16, activation='relu')
        self.fc2 = tf.keras.layers.Dense(num_classes, activation='softmax')

    def call(self, x):
        x = tf.transpose(x, perm=[128, 3, 1])
        x = tf.expand_dims(x, axis=3)
        x = self.conv1(x)
        x = self.conv2(x)
        # x = tf.squeeze(x, axis=2)
        # x = tf.transpose(x, perm=[0, 2, 1])
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.fc2(x)

        return x